
package edu.utah.seq.parsers;

import java.io.*;
import java.util.regex.*;
import edu.utah.seq.data.*;
import util.gen.*;
import java.util.*;
import util.bio.annotation.*;

/**Parses a Novoalign alignment txt file into point chromPointData, split by chromosome and strand.
 * For each sequence a single hit is assigned to the center position of the read.  Files are saved using the bar format.
 * Final positions are in interbase coordinates (0 start, end excluded).
 * Recommend: novoalign -r0.015 -q5 -d database -f pbrFiles"_prb.txt.gz" | grep chr >> results
 * If a chrAdapter.fasta was included will remove any reads that have hits to it.
 * @author david.nix@hci.utah.edu 
 **/
public class NovoalignParser {
	//fields
	private File[] dataFiles;
	private File saveDirectory;
	private File bedDirectory;
	private File pointDataDirectory;
	private File tempDirectory;
	private File workingFile;
	private String versionedGenome;
	private HashMap <String, DataOutputStream> chromOut = new HashMap <String, DataOutputStream>();
	private Pattern tab = Pattern.compile("\\t");
	private ArrayList<File> tempChrData = new ArrayList<File>();
	private float probThreshold = 13;
	private float alignmentThreshold = 60;
	private int numberPassingAlignments;
	private int numberAdapterReads = 0;
	private int numberNonAdapterReads = 0;
	private String chromosomePrefix = ">chr";
	private boolean skipUnderScoredChromosomes = true;
	private String adapterName = "chrAdapter";
	private boolean filterAdapterSequences = false;


	//constructors
	public NovoalignParser(String[] args){
		long startTime = System.currentTimeMillis();
		processArgs(args);
		System.out.println("\n"+probThreshold+"\tPosterior probability threshold");
		System.out.println(alignmentThreshold+"\tAlignment score threshold");
		System.out.println("\nParsing and filtering...");

		//for each file, parse and save to disk
		for (int i=0; i< dataFiles.length; i++){
			//set working objects and parse tag file name
			workingFile = dataFiles[i];
			System.out.print("\t"+workingFile);

			//split file to chromosome strand specific temp files
			boolean parsed;
			if (filterAdapterSequences) parsed = parseWorkingAdapterFile();
			else parsed = parseWorkingFile(); 
			if (parsed == false) Misc.printExit("\n\tError: failed to parse, aborting.\n");
			System.out.println();
		}

		//close the writers
		closeWriters();

		//load, sort, make point chromPointData, and save
		System.out.print("Loading, sorting and saving PointData and Bed files");
		makePointDataAndBedFiles();
		System.out.println();

		//cleanup
		IO.deleteDirectory(tempDirectory);

		//stats
		System.out.println("Stats:");
		System.out.println("\t"+numberPassingAlignments+"\tAlignments passing filters");
		if (filterAdapterSequences){
			System.out.println("\t"+numberAdapterReads+"\tAdapter aligning reads");
			System.out.println("\t"+numberNonAdapterReads+"\tNon Adapter aligning reads");
		}

		//finish and calc run time
		double diffTime = ((double)(System.currentTimeMillis() -startTime))/1000;
		System.out.println("\nDone! "+Math.round(diffTime)+" seconds\n");
	}

	/**Closes writers.*/
	public void closeWriters(){
		try{
			Iterator<DataOutputStream> it = chromOut.values().iterator();
			while (it.hasNext()) it.next().close();
		} catch (Exception e){
			e.printStackTrace();
		}
	}

	public void makePointDataAndBedFiles(){
		//for each composite data file
		for (int i=0; i< tempChrData.size(); i++){
			File chromDataFile = tempChrData.get(i);
			//get StartStopScores and sort
			NamedScoredStartStop[] sss = NamedScoredStartStop.loadBinary(chromDataFile);
			Arrays.sort(sss);
			//parse strand and chrom
			String name = chromDataFile.getName();
			String strand = name.substring(name.length()-1);
			String chrom = name.substring(0, name.length()-1);
			//make files
			saveBedFile(sss,chrom,strand);
			savePointData(sss,chrom,strand);

		}
	}


	/**Makes the PointData, sorts and saves.*/
	public void savePointData(NamedScoredStartStop[] sss, String chrom, String strand){
		//make Point[]
		Point[] points = new Point[sss.length];
		long totalReadLengthSize = 0;
		for (int i=0; i<sss.length; i++){
			//add to length
			totalReadLengthSize += sss[i].getLength();
			//make point
			points[i] = new Point(sss[i].getMiddle(), sss[i].getScore());
		}
		//average read length
		int averageReadLength = (int)(totalReadLengthSize/points.length);
		//make notes
		HashMap <String,String> notes = new HashMap <String,String> ();
		notes.put(BarParser.GRAPH_TYPE_TAG, BarParser.GRAPH_TYPE_BAR);
		notes.put(BarParser.SOURCE_TAG, IO.concatinateFileFullPathNames(dataFiles, ","));
		notes.put(BarParser.STRAND_TAG, strand);
		notes.put(BarParser.READ_LENGTH_TAG, +averageReadLength+"");
		notes.put(BarParser.UNIT_TAG, "Probability of originating from the given location.");
		notes.put(BarParser.DESCRIPTION_TAG, "Generated by running the NovoalignParser on Novoalign alignment file(s), the position is assigned to the middle of the read, interbase coordinates. The posterior probability thresholds was set to "+probThreshold);
		//make an Info object  public Info (String name, String versionedGenome, String chromosome, String strand, int readLength, HashMap<String,String> notes){
		Info info = new Info(chrom+strand, versionedGenome, chrom, strand, averageReadLength, notes);
		//make pd
		PointData pd = Point.extractPositionScores(points);			
		pd.setInfo(info);
		//write to file
		pd.writePointData(pointDataDirectory);
		//cleanup
		points = null;
		System.out.print(".");
	}

	/**Makes and saves the data to a bed file.*/
	public void saveBedFile(NamedScoredStartStop[] sss, String chrom, String strand){
		try {
			File bedFile = new File (bedDirectory, chrom+strand+".bed");
			PrintWriter out = new PrintWriter (new FileWriter (bedFile));
			for (int i=0; i< sss.length; i++){
				out.println(chrom+ "\t"+ sss[i].getStart()+ "\t"+ sss[i].getEnd()+ "\t"+ sss[i].getName()+ "\t"+ sss[i].getScore()+ "\t"+ strand);
			}
			out.close();
			IO.zipAndDelete(bedFile);
			System.out.print(".");
		} catch (Exception e){
			e.printStackTrace();
		}
	}





	public boolean parseWorkingFile(){
		try{
			//get reader
			BufferedReader in = IO.fetchBufferedReader(workingFile);
			String line;
			String[] tokens = null;

			//find out if the miRNA scoring mode was turned on
			int counter = 0;
			int chromIndex = -1;
			int positionIndex = -1;
			int oriIndex = -1;
			while ((line = in.readLine()) !=null){
				//does it contain 'chr'?
				if (line.startsWith("#") == false){
					if (line.contains(chromosomePrefix)){
						tokens = tab.split(line);
						if (tokens[7].startsWith(chromosomePrefix)){
							chromIndex = 7;
							positionIndex = 8;
							oriIndex = 9;
							break;
						}
						else if (tokens[8].startsWith(chromosomePrefix)){
							chromIndex = 8;
							positionIndex = 9;
							oriIndex = 10;
							break;
						}
						else Misc.printExit("\nProblem identifying chromosome column? This should start with '"+chromosomePrefix+"' and reside in column 7 or 8. See -> "+line);
					}
					if (counter++ == 1000) Misc.printExit("\nProblem identifing chromosome column? No '"+chromosomePrefix+"' found in 1st 1000 lines?\n");
				}
			}

			/*read in lines
			name type read qual matchClass alignScore prob (miRNAScore) chrom position ori
			 0    1     2   3     4            5        6                 7       8     9  
			                                                     7        8       9     10
			 */
			//reset BufferedReader
			in = IO.fetchBufferedReader(workingFile);
			counter = 0;
			while ((line = in.readLine()) !=null){
				if (++counter == 25000){
					System.out.print(".");
					counter = 0;
				}
				if (line.startsWith("#")) continue;
				tokens = tab.split(line);
				if (tokens.length < 9) continue;
				//check probability threshold
				float probScore = Float.parseFloat(tokens[6]);
				if (probScore < probThreshold) continue;
				//parse chromosome
				String chromosome = tokens[chromIndex].substring(1);
				//any underscores designating this as a splice junction match?
				if (skipUnderScoredChromosomes && chromosome.contains("_")) continue;
				//increment counter
				numberPassingAlignments++;
				//parse strand
				String strand = "+";
				if (tokens[oriIndex].equals("R")) strand = "-";
				String chrStrand = chromosome+strand;

				//calc start stop and score
				int size = tokens[2].length();
				int start = Integer.parseInt(tokens[positionIndex]) -1;
				int stop = start + size;
				float convertedScore = (float)(1- Num.antiNeg10log10(probScore));

				//get PrintWriter
				DataOutputStream dos;
				if (chromOut.containsKey(chrStrand)) dos = chromOut.get(chrStrand);
				else {
					File f = new File(tempDirectory, chrStrand);
					tempChrData.add(f);
					dos = new DataOutputStream(new FileOutputStream(f));
					chromOut.put(chrStrand, dos);
				}
				//save data
				dos.writeInt(start);
				dos.writeInt(stop);
				dos.writeFloat(convertedScore);
				dos.writeInt(size);
				dos.writeBytes(tokens[2]);
			}
			in.close();
			return true;
		} catch (Exception e){
			e.printStackTrace();
			return false;
		}
	}

	/**Reads in blocks of reads return null if at the end of the file or an ArrayList<String[]>.*/
	public ArrayList<String[]> collectReads(BufferedReader in) throws IOException{
		ArrayList<String[]> reads = new ArrayList<String[]>();
		String line;
		String readName = null;
		boolean adapterFound = false;		
		while ((line = in.readLine()) !=null){
			if (line.startsWith("#")) continue;
			String[] tokens = tab.split(line);
			if (tokens.length < 9) continue;
			//first line?
			if (readName == null) {
				readName = tokens[0];
				reads.add(tokens);
				in.mark(100000);
				if (line.contains(adapterName)) adapterFound = true;
			}
			//nope new line, is it the same read?
			else if (readName.equals(tokens[0])){
				reads.add(tokens);
				in.mark(100000);
				if (line.contains(adapterName)) adapterFound = true;
			}
			//different read 
			else {
				in.reset();
				break;
			}
		}
		//anything parsed?
		if (readName == null) return null;
		//any adapters found? if so then look for next batch
		if (adapterFound) {
			numberAdapterReads++;
			return collectReads (in);
		}
		numberNonAdapterReads++;
		return reads;
	}

	/**Parses a novoalign file that may contain hits to an adapter chromosome.  If found all alignments associated with the read are tossed.*/
	public boolean parseWorkingAdapterFile(){
		String[] t = null;
		try{
			//get reader
			BufferedReader in = IO.fetchBufferedReader(workingFile);
			String line;
			numberPassingAlignments = 0;

			//find out if the miRNA scoring mode was turned on
			int counter = 0;
			int chromIndex = -1;
			int positionIndex = -1;
			int oriIndex = -1;
			while ((line = in.readLine()) !=null){
				//does it contain 'chr'?
				if (line.startsWith("#") == false){
					if (line.contains(chromosomePrefix)){
						String[] tokens = tab.split(line);
						if (tokens[7].startsWith(chromosomePrefix)){
							chromIndex = 7;
							positionIndex = 8;
							oriIndex = 9;
							break;
						}
						else if (tokens[8].startsWith(chromosomePrefix)){
							chromIndex = 8;
							positionIndex = 9;
							oriIndex = 10;
							break;
						}
						else Misc.printExit("\nProblem identifying chromosome column? This should start with '"+chromosomePrefix+"' and reside in column 7 or 8. See -> "+line);
					}
					if (counter++ == 1000) Misc.printExit("\nProblem identifing chromosome column? No '"+chromosomePrefix+"' found in 1st 1000 lines?\n");
				}
			}

			/*read in lines
			name type read qual matchClass alignScore prob (miRNAScore) chrom position ori
			 0    1     2   3     4            5        6                 7       8     9  
			                                                     7        8       9     10
			 */
			//reset BufferedReader
			in = IO.fetchBufferedReader(workingFile);
			ArrayList<String[]> reads;
			counter = 0;
			while ((reads = collectReads(in)) !=null){
				int numReads = reads.size();
				if (++counter == 25000){
					System.out.print(".");
					counter = 0;
				}
				for (int i=0; i< numReads; i++){
					t = reads.get(i);
					//check probability threshold
					float probScore = Float.parseFloat(t[6]);
					if (probScore < probThreshold) continue;
					//check alignment threshold
					float alignmentScore = Float.parseFloat(t[5]);
					if (alignmentScore > alignmentThreshold) continue;
					//parse chromosome
					String chromosome = t[chromIndex].substring(1);
					//any underscores designating this as a splice junction match?
					if (skipUnderScoredChromosomes && chromosome.contains("_")) continue;
					//increment counter
					numberPassingAlignments++;
					
					//parse strand
					String strand = "+";
					if (t[oriIndex].equals("R")) strand = "-";
					String chrStrand = chromosome+strand;

					//calc start stop and score
					int size = t[2].length();
					int start = Integer.parseInt(t[positionIndex]) -1;
					int stop = start + size;
					float convertedScore = (float)(1- Num.antiNeg10log10(probScore));

					//get PrintWriter
					DataOutputStream dos;
					if (chromOut.containsKey(chrStrand)) dos = chromOut.get(chrStrand);
					else {
						File f = new File(tempDirectory, chrStrand);
						tempChrData.add(f);
						dos = new DataOutputStream(new FileOutputStream(f));
						chromOut.put(chrStrand, dos);
					}
					//save data
					dos.writeInt(start);
					dos.writeInt(stop);
					dos.writeFloat(convertedScore);
					//save sequence
					dos.writeInt(size);
					dos.writeBytes(t[2]);
				}
			}
			in.close();
			return true;
		} catch (Exception e){
			e.printStackTrace();
			System.out.println("Bad line\n"+Misc.stringArrayToString(t, "\t"));
			return false;
		}
	}


	public static void main(String[] args) {
		if (args.length ==0){
			printDocs();
			System.exit(0);
		}
		new NovoalignParser(args);
	}		

	/**This method will process each argument and assign new varibles*/
	public void processArgs(String[] args){
		Pattern pat = Pattern.compile("-[a-z]");
		File forExtraction = null;
		System.out.println("\nArguments: "+Misc.stringArrayToString(args, " ")+"\n");
		for (int i = 0; i<args.length; i++){
			String lcArg = args[i].toLowerCase();
			Matcher mat = pat.matcher(lcArg);
			if (mat.matches()){
				char test = args[i].charAt(1);
				try{
					switch (test){
					case 'f': forExtraction = new File(args[i+1]); i++; break;
					case 'v': versionedGenome = args[i+1]; i++; break;
					case 'r': saveDirectory = new File (args[i+1]); i++; break;
					case 'p': probThreshold = Float.parseFloat(args[++i]); break;
					case 'q': alignmentThreshold = Float.parseFloat(args[++i]); break;
					case 'c': chromosomePrefix = args[++i]; break;
					case 'a': filterAdapterSequences = true; break;
					case 'n': adapterName = args[++i]; break;
					case 'h': printDocs(); System.exit(0);
					default: System.out.println("\nProblem, unknown option! " + mat.group());
					}
				}
				catch (Exception e){
					Misc.printExit("\nSorry, something doesn't look right with this parameter: -"+test+"\n");
				}
			}
		}
		//pull files
		File[][] tot = new File[3][];
		tot[0] = IO.extractFiles(forExtraction,".txt");
		tot[1] = IO.extractFiles(forExtraction,".txt.zip");
		tot[2] = IO.extractFiles(forExtraction,".txt.gz");

		dataFiles = IO.collapseFileArray(tot);
		if (dataFiles == null || dataFiles.length==0) dataFiles = IO.extractFiles(forExtraction);
		if (dataFiles == null || dataFiles.length ==0 || dataFiles[0].canRead() == false) Misc.printExit("\nError: cannot find your xxx.txt(.zip/.gz) file(s)!\n");
		if (versionedGenome == null) Misc.printExit("\nPlease enter a genome version recognized by UCSC, see http://genome.ucsc.edu/FAQ/FAQreleases.\n");
		if (saveDirectory == null)  Misc.printExit("\nPlease provide a directory in which to save the parsed data.\n");
		else if (saveDirectory.exists() == false) saveDirectory.mkdir();
		tempDirectory = new File (saveDirectory,"TempFilesDelme");
		tempDirectory.mkdir();
		bedDirectory = new File (saveDirectory, "Bed_"+versionedGenome);
		bedDirectory.mkdir();
		pointDataDirectory = new File (saveDirectory, "PointData");
		pointDataDirectory.mkdir();
	}	

	public static void printDocs(){
		System.out.println("\n" +
				"**************************************************************************************\n" +
				"**                            Novoalign Parser: March 2009                          **\n" +
				"**************************************************************************************\n" +
				"Parses Novoalign xxx.txt(.zip/.gz) files into center position binary PointData xxx.bar\n" +
				"files and xxx.bed files. Interbase coordiantes (zero based, end excluded).\n" +
				"These can be directly viewed in the Integrated Genome Browser. We recommend using the\n" +
				"following settings while running Novoalign 'novoalign -r0.2 -q5 -d yourDataBase -f\n"+
				"your_prb.txt | grep '>chr' > yourResultsFile.txt'.\n" +

				"\nOptions:\n"+
				"-v Versioned Genome (ie H_sapiens_Mar_2006), see UCSC Browser,\n"+
				"      http://genome.ucsc.edu/FAQ/FAQreleases.\n" +
				"-f The full path directory/file name of your Novoalign xxx.txt(.zip or .gz) file(s).\n" +
				"-r Full path directory name for saving the results.\n"+
				"-p Posterior probability threshold (-10Log10(prob)) of being incorrect, defaults to 13\n"+
				"      (0.05). Larger numbers are more stringent. The parsed scores are delogged and\n" +
				"      converted to 1-prob.\n"+
				"-q Alignment score threshold, smaller numbers are more stringent, defaults to 60\n"+
				"-c Chromosome prefix, defaults to '>chr'.\n"+
				"-a Remove reads where one of their alignments align to an adapter chromosome. Assumes\n" +
				"      that an artificial chrAdapter.fasta was included while running the Novoaligner.\n"+
				"-n Name of the adapter chromosome, defaults to 'chrAdapter'\n"+


				"\nExample: java -Xmx1500M -jar pathToUSeq/Apps/NovoalignParser -f /Novo/Run7/\n" +
				"     -v H_sapiens_Mar_2006 -p 20 -q 30 -r /Novo/Run7/Parsed_13/ \n\n" +

		"**************************************************************************************\n");

	}	

}
